from myglobal import *

from utils.loc import *
from pylab import *
def plot_stations_loc(stations,names2show=[], points2show=[], 
                      KM=True, special_names=[], c='',
                      map_style = 2,bounds=None,zoom=14,ax=None, dx=0.02,dy=0.01,
                      filename=None,
                      ) -> None:
    '''
    绘制地震台站位置
    stations : 台站位置信息
    names2show : 显示的台站名称
    points2show : 显示的点名称
    KM : 是否以KM为单位显示
    special_names : 特殊台站名称
    map_style : 地图样式
    bounds : 地图边界
    zoom : 地图缩放比例
    filename : 保存图片名称
    map_style : 地图样式
    ax : 绘制的轴
    dx : 经度间隔
    dy : 纬度间隔
    filename : 保存图片名称
    '''
    rcParams['font.family'] = 'Arial'
    rcParams['font.size'] = '10'

    if ax is None:
        fig = plt.figure(figsize=[5,4], dpi=600)
        ax = plt.subplot(111)

    if bounds is None:
        bounds = [115.16,38.57,115.36,38.68]
        # bounds = [115.22,38.60,115.26,38.63]
        # bounds = [115.66,39.06,115.92,39.91]
        # bounds = [115.77,39.00,115.92,39.14]
    #创建图框
    plt.sca(ax)
    #添加地图底图
    tbd.plot_map(plt,bounds,zoom = zoom,style = map_style)
    tbd.plotscale(ax,bounds = bounds,textsize = 8,compasssize = 1,accuracy = 500,unit='km',rect = [0.6,0.06],zorder = 15, style=2)
    color='k' if map_style==2 else 'red'
    lats, lons = [],[]
    for stat_name in stations.keys():
        lat, lon= stations[stat_name][:2]
        
        if stat_name in names2show:
            # plt.text(lon, lat-0.001, stat_name, color='tab:red', fontdict={'fontsize':7}, rotation=0)
            plt.text(lon, lat-0.0045, f'R{points2show.index(stat_name)+1:02d}', color='k', fontdict={'fontsize':8}, rotation=0)
            # plt.scatter(lon, lat,marker='^', s=30, color='tab:red', edgecolor='none')
        if stat_name in points2show:
            plt.scatter(lon, lat,marker='^', s=30, color='blue', edgecolor='none')
        if stat_name in names2show or stat_name in points2show:
            lats.append(lat)
            lons.append(lon)
    
    # xs,xe = int(min(lons)/dx)*dx,int(max(lons)/dx)*dx
    # ys,ye = int(min(lats)/dy)*dy,int(max(lats)/dy+1)*dy
    ys,ye = bounds[1], bounds[3]
    xs,xe = bounds[0], bounds[2]
    print(xs,xe,ys,ye)
    xticks = np.arange(xs,xe,dx)
    xtickslabel = ['{:3.2f}'.format(i) for i in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xtickslabel)
    yticks = np.arange(ys,ye+dy,dy)
    ytickslabel = ['{:2.2f}'.format(i) for i in yticks]
    ax.set_yticks(yticks)
    ax.set_yticklabels(ytickslabel)
    

    Y1, X1 =  [38+35/60+39.65/3600,115+12/60+26.73/3600]
    Y2, X2 =  [38+36/60+3.62/3600,115+13/60+11.58/3600]
    Y3, X3 =  [38+39/60+31.17/3600,115+19/60+41.57/3600]
    Y4, X4 =  [38.6731216,115.3556417]
    Y5, X5 =  [38+42/60,115+24/60+22.25/3600]
    
    ax.plot([X1,X2],[Y1,Y2],'k-',lw=2,zorder=0)
    ax.plot([X2,X3],[Y2,Y3],'g-',lw=2,zorder=0)
    ax.plot([X3,X4],[Y3,Y4],'k-',lw=2,zorder=0)
    ax.plot([X4,X5],[Y4,Y5],'g-',lw=2,zorder=0)

    ax.set_xlim([xs,xe])
    ax.set_ylim([ys,ye])
    
    if filename is not None:
        plt.xlabel('longitude')
        plt.ylabel('latitude')
        plt.tight_layout()
        plt.savefig(filename)
        plt.savefig(filename+'.pdf')
    return ax

s_info = load_loc('./loc/loc_all_from_log_2303.csv')
file_name = './figures/loc.png'
points2show = list(reversed(LINES['HSR_SHORT']))+\
              ['DZ154', 'DZ155', 'DZ156', 'DZ157', 'DZ158', 'DZ159', 'DZ160']+\
              LINES['HF']+LINES['VF']    
names2show = [
    'DZ155', 'DZ156', 'DZ157', 'DZ158', 'DZ159',
    'DZ090','DZ064',
    # 'DZ132','DZ146',
    # 'DZ117','DZ131'
]
plot_stations_loc(s_info,names2show=names2show,points2show=points2show,filename=file_name,map_style=11,
                  dx=0.04,dy=0.04,
                  bounds=[115.21,38.58,115.37,38.70],
                  c='tab:blue')